# 验试使用KAN做数值回归入口程序
import argparse
from typing import Dict
import torch
from kan import KAN
from apps.wfs.wfs_config import WfsConfig as WG

class TcrKanApp(object):
    def __init__(self):
        self.name = 'apps.wfs.tcr_kan_app.TcrKanApp'


        '''
        
        self.l1_fc = nn.Linear(4, 32, dtype=torch.float64)
        self.l1_relu = nn.ReLU()
        self.l2_fc = nn.Linear(32, 128, dtype=torch.float64)
        self.l2_relu = nn.ReLU()
        self.l3_fc = nn.Linear(128, 512, dtype=torch.float64)
        self.l3_relu = nn.ReLU()
        self.l4_fc = nn.Linear(512, 128, dtype=torch.float64)
        self.l4_relu = nn.ReLU()
        self.l5_fc = nn.Linear(128, 32, dtype=torch.float64)
        self.l5_relu = nn.ReLU()
        self.l6_fc = nn.Linear(32, N, dtype=torch.float64)
        self.y_fn = nn.Sigmoid()
        '''

    @staticmethod
    def startup(params:Dict = {}) -> None:
        print(f'KAN回归 v0.0.1')
        torch.set_default_dtype(torch.float64)
        # create a KAN: 2D inputs, 1D output, and 5 hidden neurons. cubic spline (k=3), 5 grid intervals (grid=5).
        model = KAN(width=[4,32,128,512,128,32,WG.N], grid=100, k=3, seed=42, device=WG.device)
        # 生成数据集
        X = torch.tensor([[WG.N, WG.f, WG.d, WG.theta]], dtype=torch.float64)
        dataset = {}
        dataset['train_input'] = torch.tensor([[WG.N, WG.f, WG.d, WG.theta]], dtype=torch.float64).to(WG.device)
        dataset['test_input'] = torch.tensor([[WG.N, WG.f, WG.d, WG.theta]], dtype=torch.float64).to(WG.device)
        dataset['train_label'] = torch.tensor([[0.26221649, 0.51874705, 0.81196007, 1., 1., 0.81196007, 0.51874705, 0.26221649]], dtype=torch.float64).to(WG.device)
        dataset['test_label'] = torch.tensor([[0.26221649, 0.51874705, 0.81196007, 1., 1., 0.81196007, 0.51874705, 0.26221649]], dtype=torch.float64).to(WG.device)
        # plot KAN at initialization
        model(dataset['train_input'])
        # model.plot()
        # train the model
        model.fit(dataset, opt="LBFGS", steps=50000, lamb=0.001)
        # model = KAN.loadckpt(path='./model')
        print(f'##### ^_^ The End! ^_^ #####')


def main(params:Dict = {}) -> None:
    TcrKanApp.startup(params=params)

def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--run_mode', action='store',
        type=int, default=1, dest='run_mode',
        help='run mode'
    )
    return parser.parse_args()

if '__main__' == __name__:
    args = parse_args()
    params = vars(args)
    main(params=params)